# coding=utf-8
import cv2 as cv
import numpy as np
import glob
np.set_printoptions(precision=6, suppress=True)
criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)
fname = glob.glob("E:/Study/CSDN人工智能/1-week-vision/作业/chessboard/*.jpg")
w = 9
h = 6
objp = np.zeros((w*h,3),np.float32)
objp[:,:2] = np.mgrid[0:w,0:h].T.reshape(-1,2)
objpoints = []
imgpoints = []
for fn in fname:
    img = cv.resize(cv.imdecode(np.fromfile(fn,dtype=np.uint8),-1),dsize=(0,0),fx=0.5,fy=0.5)
    gray = cv.cvtColor(img,cv.COLOR_BGR2GRAY)
    ret, corners = cv.findChessboardCorners(gray,(w,h))
    if ret:
        cv.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria)
        objpoints.append(objp)
        imgpoints.append(corners)
        cv.drawChessboardCorners(img,(w,h),corners,ret)
        cv.imshow('findCorners',img)
        cv.waitKey(1)
cv.destroyAllWindows()

img = cv.resize(cv.imdecode(np.fromfile(fname[0], dtype=np.uint8), -1), dsize=(0, 0), fx=0.5, fy=0.5)
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
ret,mtx,dist,rvecs,tvecs = cv.calibrateCamera(objpoints,imgpoints,gray.shape[::-1],None,None)
print(mtx)
print(dist)
print(rvecs[0])
print(tvecs[0])

total_error = 0
for i in range(len(objpoints)):
    imgpoints2, _ = cv.projectPoints(objpoints[i], rvecs[i], tvecs[i], mtx, dist)
    error = cv.norm(imgpoints[i],imgpoints2, cv.NORM_L2)/len(imgpoints2)
    total_error += error
print("total error: ", total_error/len(objpoints))

fl = "E:/Study/CSDN人工智能/2-week-vision/作业/chessboard/0.jpg"
fr = "E:/Study/CSDN人工智能/2-week-vision/作业/chessboard/1.jpg"

img = cv.resize(cv.imdecode(np.fromfile(fr,dtype=np.uint8),-1),dsize=(0,0),fx=0.5,fy=0.5)
gray = cv.cvtColor(img,cv.COLOR_BGR2GRAY)
ret, rpoint = cv.findChessboardCorners(gray,(w,h))
rpoint=cv.cornerSubPix(gray,rpoint,(11,11),(-1,-1),criteria)

img = cv.resize(cv.imdecode(np.fromfile(fl,dtype=np.uint8),-1),dsize=(0,0),fx=0.5,fy=0.5)
gray = cv.cvtColor(img,cv.COLOR_BGR2GRAY)
ret, lpoint = cv.findChessboardCorners(gray,(w,h))
lpoint=cv.cornerSubPix(gray,lpoint,(11,11),(-1,-1),criteria)

cv.destroyAllWindows()

E=cv.findEssentialMat(rpoint,lpoint,mtx)[0]
_,R,t,_=cv.recoverPose(E,rpoint,lpoint,mtx)
print("E: \n",E)
print("R: \n",R)
print("t: \n",t)

theta_y=np.arcsin(-R[0,2])
theta_z=np.arccos(R[0,0]/np.cos(theta_y))
theta_x=np.arccos(R[2,2]/np.cos(theta_y))

print("theta_x: ",theta_x/np.pi*180)
print("theta_y: ",theta_y/np.pi*180)
print("theta_z: ",theta_z/np.pi*180)